#include "master_server.h"

rapidjson::Document ServerInfo::toJson() const
{
    rapidjson::Document json;
    json.SetObject();
    rapidjson::Document::AllocatorType& allocator = json.GetAllocator();

    rapidjson::Value urlValue;
    urlValue.SetString(url.c_str(),allocator);

    rapidjson::Value roleValue;
    roleValue.SetInt(static_cast<int>(role));

    json.AddMember("url",urlValue,allocator);
    json.AddMember("role",roleValue,allocator);
    return json;
}

ServerInfo ServerInfo::fromJson(const rapidjson::Document &value)
{
    ServerInfo info;
    info.url = value["url"].GetString();
    info.role = static_cast<ServerRole>(value["role"].GetInt());
    return info;
}

MasterServer::MasterServer(const std::string &etcdEndPoint, int httpPort):etcdClient_(etcdEndPoint),httpPort_(httpPort)
{   
    httpServer_.Get(
        "/getNodeInfo",[this](const httplib::Request &req, httplib::Response &res){
            getNodeInfo(req,res);
        }
    );

    httpServer_.Post(
        "/addNode",[this](const httplib::Request &req, httplib::Response &res){
            addNode(req,res);
    });

    httpServer_.Delete(
        "/removeNode",[this](const httplib::Request &req, httplib::Response &res){
            removeNode(req,res);
    });

    httpServer_.Get(
        "/getInstance",[this](const httplib::Request &req, httplib::Response &res){
            getInstance(req,res);
        }
    );
    httpServer_.Get("/getPartitionConfig", [this](const httplib::Request& req, httplib::Response& res) {
        getPartitionConfig(req, res);
    });
    httpServer_.Post("/updatePartitionConfig", [this](const httplib::Request& req, httplib::Response& res) {
        updatePartitionConfig(req, res);
    });
    httpServer_.Post("/updateUserCredentials", [this](const httplib::Request& req, httplib::Response& res) {
        updateUserCredentials(req, res);
    });
    httpServer_.Get("/getUserCredentials", [this](const httplib::Request& req, httplib::Response& res) {
        getUserCredentials(req, res);
    });


    // 启动定时器
    startNodeUpdateTimer();
}

void MasterServer::run(){
    httpServer_.listen("0.0.0.0", httpPort_);
}

void MasterServer::startNodeUpdateTimer() {
    std::thread([this]() {
        while (true) { // 这里可能需要一种更优雅的退出机制
            std::this_thread::sleep_for(std::chrono::seconds(10)); // 每10秒更新一次
            updateNodesStatus();
        }
    }).detach();
}

void MasterServer::updateUserCredentials(const httplib::Request& req, httplib::Response& res) {
    rapidjson::Document doc;
    doc.Parse(req.body.c_str());

    if (!doc.IsObject() || !doc.HasMember("instanceId") || !doc.HasMember("credentials")) {
        setResponse(res, 1, "Invalid JSON format");
        return;
    }

    std::string instanceId = doc["instanceId"].GetString();
    std::string etcdKey = "/instancesConfig/" + instanceId + "/credentials";

    try {
        etcd::Response etcdResponse = etcdClient_.get(etcdKey).get();
        rapidjson::Document currentDoc;
        if (etcdResponse.is_ok() && !etcdResponse.value().as_string().empty()) {
            currentDoc.Parse(etcdResponse.value().as_string().c_str());
        } else {
            currentDoc.SetObject(); // 初始化一个空对象
        }

        const rapidjson::Value& newCredentials = doc["credentials"];
        rapidjson::Value::MemberIterator itr = currentDoc.FindMember(newCredentials["username"].GetString());

        if (itr != currentDoc.MemberEnd()) {
            // 更新现有用户的密码
            itr->value["password"].SetString(newCredentials["password"].GetString(), currentDoc.GetAllocator());
        } else {
            // 添加新用户
            rapidjson::Value key(newCredentials["username"].GetString(), currentDoc.GetAllocator());
            rapidjson::Value value(rapidjson::kObjectType);

            // 创建新的 rapidjson::Value 对象来存储用户名和密码
            rapidjson::Value username(newCredentials["username"].GetString(), currentDoc.GetAllocator());
            rapidjson::Value password(newCredentials["password"].GetString(), currentDoc.GetAllocator());

            value.AddMember("username", username, currentDoc.GetAllocator());
            value.AddMember("password", password, currentDoc.GetAllocator());

            currentDoc.AddMember(key, value, currentDoc.GetAllocator());
        }

        // 将更新后的数据保存回 etcd
        rapidjson::StringBuffer buffer;
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
        currentDoc.Accept(writer);

        etcdClient_.set(etcdKey, buffer.GetString()).get();
        setResponse(res, 0, "User credentials updated successfully");
    } catch (const std::exception& e) {
        setResponse(res, 1, std::string("Error updating user credentials: ") + e.what());
    }
}
void MasterServer::setResponse(httplib::Response& res,int retcode,const std::string& message,const rapidjson::Document* data){
    rapidjson::Document json_response;
    json_response.SetObject();
    rapidjson::Document::AllocatorType& allocator = json_response.GetAllocator();

    json_response.AddMember("retcode", retcode, allocator);
    json_response.AddMember("message", rapidjson::Value(message.c_str(),allocator), allocator);

    if(data!=nullptr&&data->IsObject()){
        rapidjson::Value dataValue(rapidjson::kObjectType);
        dataValue.CopyFrom(*data, allocator);
        json_response.AddMember("data", dataValue, allocator);
    }

    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
    json_response.Accept(writer);

    res.set_content(buffer.GetString(),"application/json");
}

void MasterServer::getNodeInfo(const httplib::Request &req, httplib::Response &res)
{
    auto instanceId=req.get_param_value("instanceId");
    auto nodeId=req.get_param_value("nodeId");
    try{
        std::string etcdKey="/instances/"+instanceId+"/nodes/"+nodeId;
        etcd::Response etcdResponse=etcdClient_.get(etcdKey).get();

        if(!etcdResponse.is_ok()){
            setResponse(res,1,"Error accessing etcd:" + etcdResponse.error_message());
            return;
        }

        rapidjson::Document json_data;
        json_data.SetObject();
        rapidjson::Document::AllocatorType& allocator = json_data.GetAllocator();

        // 解析节点地址
        rapidjson::Document nodeDoc;
        nodeDoc.Parse(etcdResponse.value().as_string().c_str());
        if(!nodeDoc.IsObject()){
            setResponse(res,1,"Invalid JSON format");
            return;
        }

        json_data.AddMember("instanceId",rapidjson::Value(instanceId.c_str(),allocator),allocator);
        json_data.AddMember("nodeId",rapidjson::Value(nodeId.c_str(),allocator),allocator);
        json_data.AddMember("nodeInfo",nodeDoc,allocator);

        setResponse(res,0,"Node info retrieved successfully",&json_data);
    }catch(const std::exception& e) {
        setResponse(res,1,"Failed to retrieve node info:"+std::string(e.what()) );
    }
}

void MasterServer::addNode(const httplib::Request &req, httplib::Response &res)
{
    rapidjson::Document json_request;
    json_request.Parse(req.body.c_str());
    if(!json_request.IsObject()){
        setResponse(res,1,"Invalid JSON format");
        return;
    }

    try
    {
        std::string instanceId=json_request["instanceId"].GetString();
        std::string nodeId=json_request["nodeId"].GetString();
        std::string etcdKey="/instance/"+instanceId+"/nodes/"+nodeId;

        rapidjson::StringBuffer buffer;
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
        json_request.Accept(writer);

        etcdClient_.set(etcdKey,buffer.GetString()).get();
        setResponse(res,0,"Node added successfully");
    }
    catch(const std::exception& e)
    {
        setResponse(res,1,"Add node err:"+  std::string(e.what()));
    }

    
}

void MasterServer::removeNode(const httplib::Request &req, httplib::Response &res)
{
    auto instanceId=req.get_param_value("instanceId");
    auto nodeId=req.get_param_value("nodeId");
    std::string etcdKey="/instance/"+instanceId+"/nodes/"+nodeId;

    try{
        etcd::Response etcdResponse = etcdClient_.rm(etcdKey).get();
        if(!etcdResponse.is_ok()){
            setResponse(res,1,"Failed to remove node:"+etcdResponse.error_message());
            return;
        }
        setResponse(res,0,"Node removed successfully");
    }catch(std::exception &e){
        setResponse(res,1,"Failed to remove node:"+std::string(e.what()));
        return;
    }
}

void MasterServer::getInstance(const httplib::Request &req, httplib::Response &res)
{
    auto instanceId=req.get_param_value("instanceId");
    std::string etcdPrefix="/instance/"+instanceId+"/nodes/";
    GlobalLogger->info("Getting instance information for instanceId: {}",instanceId);

    try
    {
        GlobalLogger->debug("etcdPrefix:{}",etcdPrefix);
        etcd::Response r=etcdClient_.ls(etcdPrefix).get();
        GlobalLogger->debug("etcd ls response received");

        if(!r.is_ok()){
            GlobalLogger->error("Error accessing etcd: " + r.error_message());
            setResponse(res,1,"Error accessing etcd: " + r.error_message());
            return;
        }

        const auto& keys=r.keys();
        const auto& values=r.values();

        rapidjson::Document json;
        json.SetObject();
        rapidjson::Document::AllocatorType& allocator=json.GetAllocator();

        rapidjson::Value nodeArray(rapidjson::kArrayType);

        for(size_t i=0;i<keys.size();i++){
            GlobalLogger->debug("Processing key: {}" , keys[i]);
            rapidjson::Document doc;
            doc.Parse(values[i].as_string().c_str());
            if(!doc.IsObject()){
                GlobalLogger->error("Invalid JSON format for key: {}", keys[i]);
                continue;
            }
            rapidjson::Value node(doc,allocator);
            nodeArray.PushBack(node,allocator);
        }

        json.AddMember("instanceId",rapidjson::Value(instanceId.c_str(),allocator),allocator);
        json.AddMember("nodes",nodeArray,allocator);

        GlobalLogger->info("Instance info retrieved successfully for instanceId: {}", instanceId);
        setResponse(res,0,"Instance info retrieved successfully",&json);
    }
    catch(const std::exception& e)  
    {
        GlobalLogger->error("Exception accessing etcd:{}",e.what());
        setResponse(res,1,"Exception accessing etcd"+std::string(e.what()));
    }
}

void MasterServer::getPartitionConfig(const httplib::Request& req, httplib::Response& res) {
    auto instanceId = req.get_param_value("instanceId");
    try {
        PartitionConfig config = doGetPartitionConfig(instanceId);

        rapidjson::Document doc;
        doc.SetObject();
        rapidjson::Document::AllocatorType& allocator = doc.GetAllocator();

        // 添加分区配置信息到响应
        doc.AddMember("partitionKey", rapidjson::Value(config.partitionKey.c_str(), allocator), allocator);
        doc.AddMember("numberOfPartitions", config.numberOfPartitions, allocator);

        rapidjson::Value partitionsArray(rapidjson::kArrayType);
        for (const auto& partition : config.partitions) {
            rapidjson::Value partitionObj(rapidjson::kObjectType);
            partitionObj.AddMember("partitionId", partition.partitionId, allocator);
            partitionObj.AddMember("nodeId", rapidjson::Value(partition.nodeId.c_str(), allocator), allocator);
            partitionsArray.PushBack(partitionObj, allocator);
        }
        doc.AddMember("partitions", partitionsArray, allocator);

        setResponse(res, 0, "Partition config retrieved successfully", &doc);
    } catch (const std::exception& e) {
        setResponse(res, 1, "Exception occurred: " + std::string(e.what()));
    }
}

void MasterServer::updatePartitionConfig(const httplib::Request &req, httplib::Response &res)
{
    rapidjson::Document doc;
    doc.Parse(req.body.c_str());
    if(!doc.IsObject()){
        setResponse(res,1,"Invalid JSON format");
        return;
    }

    try{
        std::string instanceId=doc["instanceId"].GetString(); 
        std::string partitionKey=doc["partitionKey"].GetString();
        int numberOfPartitions=doc["numberOfPartitions"].GetInt();

        std::list<Partition> partitionList;
        const rapidjson::Value& partitions=doc["partitions"];
        for(const auto& partition:partitions.GetArray()){
            int partitionId=partition["partitionId"].GetInt();
            std::string nodeId=partition["nodeId"].GetString();
            partitionList.push_back({partitionId,nodeId});
        }
        doUpdatePartitionConfig(instanceId,partitionKey,numberOfPartitions,partitionList);
        setResponse(res,0,"Partition config updated successfully");
    }catch(const std::exception& e)
    {
        setResponse(res,1,"Failed to update partition config:"+std::string(e.what()));
    }
}

void MasterServer::doUpdatePartitionConfig(const std::string& instanceId,const std::string& partitionKey,int numberOfPartitions,const std::list<Partition>& partitions) {
    std::string etcdKey = "/instancesConfig/" + instanceId + "/partitionConfig";
    
    rapidjson::Document doc;
    doc.SetObject();
    rapidjson::Document::AllocatorType& allocator = doc.GetAllocator();
    
    doc.AddMember("partitionKey", rapidjson::Value(partitionKey.c_str(), allocator), allocator);
    doc.AddMember("numberOfPartitions", numberOfPartitions, allocator);
    
    rapidjson::Value partitionsArray(rapidjson::kArrayType);
    for (const auto& partition : partitions) {
        rapidjson::Value partitionObj(rapidjson::kObjectType);
        partitionObj.AddMember("partitionId", partition.partitionId, allocator);
        partitionObj.AddMember("nodeId", rapidjson::Value(partition.nodeId.c_str(), allocator), allocator);
        partitionsArray.PushBack(partitionObj, allocator);
    }
    doc.AddMember("partitions", partitionsArray, allocator);
    
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
    doc.Accept(writer);

    
    
    etcdClient_.set(etcdKey, buffer.GetString()).get();
    GlobalLogger->info("Partition config updated for instance: {}", instanceId);
}

void MasterServer::updateNodesStatus()
{
    CURL *curl=curl_easy_init();
    if(!curl){
        GlobalLogger->error("Failed to initialize curl");
        return;
    }

    try{
        std::string nodesKeyPrefix="/instances/";
        GlobalLogger->debug("Fetching nodes from etcd...");
        etcd::Response etcdResponse=etcdClient_.ls(nodesKeyPrefix).get();

        for(size_t i=0;i<etcdResponse.keys().size();++i){
            const std::string& nodeKey=etcdResponse.keys()[i];
            const std::string& nodeValue=etcdResponse.values()[i].as_string();

            rapidjson::Document nodeDoc;
            nodeDoc.Parse(nodeValue.c_str());
            if(!nodeDoc.IsObject()){
                GlobalLogger->warn("invalid JSON format for node: {}",nodeKey);
                continue;
            }

            std::string getNodeUrl=std::string(nodeDoc["url"].GetString())+"/admin/getNode";
            GlobalLogger->debug("Sending request to node: {}",getNodeUrl);

            curl_easy_setopt(curl, CURLOPT_URL, getNodeUrl.c_str());
            curl_easy_setopt(curl,CURLOPT_WRITEFUNCTION,writeCallback);
            std::string responseData;
            curl_easy_setopt(curl, CURLOPT_WRITEDATA, &responseData);

            CURLcode res = curl_easy_perform(curl);
            bool needsUpdate=false;

            if(res != CURLE_OK){
                GlobalLogger->error("curl_easy_perform() failed: {}", curl_easy_strerror(res));
                nodeErrCounts[nodeKey]++;
                if(nodeErrCounts[nodeKey]>=5&&nodeDoc["status"].GetInt()!=0){
                    nodeDoc["status"].SetInt(0);
                    needsUpdate=true;
                }
            }else{
                nodeErrCounts[nodeKey]=0;
                if(nodeDoc["status"].GetInt()!=1){
                    nodeDoc["status"].SetInt(1);
                    needsUpdate=true;
                }

                rapidjson::Document getNodeResponse;
                getNodeResponse.Parse(responseData.c_str());
                if(getNodeResponse.HasMember("node")&&getNodeResponse["node"].IsObject()){
                    std::string state=getNodeResponse["node"]["state"].GetString();
                    int newRole=(state=="leader")?1:0;

                    if(nodeDoc["role"].GetInt()!=newRole){
                        nodeDoc["role"].SetInt(newRole);
                        needsUpdate=true;
                    }
                }
            }
            GlobalLogger->debug("updated check needsUpdate: {}",needsUpdate);

            if(needsUpdate){
                rapidjson::StringBuffer buffer;
                rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
                nodeDoc.Accept(writer);

                etcdClient_.set(nodeKey,buffer.GetString()).get();
                GlobalLogger->info("Updated node status and role: {}",nodeKey); 
            }
        }
    }catch(const std::exception& e){
        GlobalLogger->error("Error updating nodes status: {}",e.what());
    }
    curl_easy_cleanup(curl);
}

size_t MasterServer::writeCallback(void *contents, size_t size, size_t nmemb, void *userp) {
    ((std::string*)userp)->append((char*)contents, size * nmemb);
    return size * nmemb;
}
void MasterServer::getUserCredentials(const httplib::Request &req, httplib::Response &res)
{
    auto instanceId=req.get_param_value("instanceId");
    std::string etcdKey=std::string("/instancesConfig/") + instanceId + "/credentials";

    try{
        etcd::Response etcdResponse=etcdClient_.get(etcdKey).get();
        if(!etcdResponse.is_ok()){ 
            setResponse(res,1,"Error accessing etcd: "+etcdResponse.error_message());
            return;
        }

        rapidjson::Document doc;
        doc.Parse(etcdResponse.value().as_string().c_str());
        if(!doc.IsObject()){
            setResponse(res,1,"Invalid JSON format in etcd");
            return;
        }
        setResponse(res,0,"user credentials retrieved successfully",&doc);
    }catch(std::exception &e){
        setResponse(res,1,"Error occurred: "+std::string(e.what()));
        return;
    }
}
PartitionConfig MasterServer::doGetPartitionConfig(const std::string& instanceId) {
    PartitionConfig config;
    std::string etcdKey = "/instancesConfig/" + instanceId + "/partitionConfig";
    etcd::Response etcdResponse = etcdClient_.get(etcdKey).get();
    rapidjson::Document doc;
    doc.Parse(etcdResponse.value().as_string().c_str());

    if (doc.IsObject()) {
        if (doc.HasMember("partitionKey")) {
            config.partitionKey = doc["partitionKey"].GetString();
        }
        if (doc.HasMember("numberOfPartitions")) {
            config.numberOfPartitions = doc["numberOfPartitions"].GetInt();
        }
        if (doc.HasMember("partitions") && doc["partitions"].IsArray()) {
            for (const auto& partition : doc["partitions"].GetArray()) {
                Partition p;
                p.partitionId = partition["partitionId"].GetInt();
                p.nodeId = partition["nodeId"].GetString();
                config.partitions.push_back(p);
            }
        }
    }

    return config;
}
